#! /usr/bin/python3
#
# memleak   Trace and display outstanding allocations to detect
#           memory leaks in user-mode processes and the kernel.
#
# USAGE: memleak [-h] [-p PID] [-t] [-a] [-o OLDER] [-c COMMAND]
#                [--combined-only] [--wa-missing-free] [-s SAMPLE_RATE]
#                [-T TOP] [-z MIN_SIZE] [-Z MAX_SIZE] [-O OBJ]
#                [interval] [count]
#
# Licensed under the Apache License, Version 2.0 (the "License")
# Copyright (C) 2016 Sasha Goldshtein.

from __future__ import annotations
from bcc import BPF
from time import sleep
from datetime import datetime
import resource
import argparse
import subprocess
import os
import sys
from typing import Dict, Tuple, Sequence, List, Any
import json
import traceback
from ctypes import *
import re
from shell_pipe import ShellPipe

p = ShellPipe()

class Allocation(object):
    def __init__(self, stack, size):
        self.stack = stack
        self.count = 1
        self.size = size

    def update(self, size):
        self.count += 1
        self.size += size

def run_command_get_output(command):
        p = subprocess.Popen(command.split(),
                stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
        return iter(p.stdout.readline, b'')

def run_command_get_pid(command):
        p = subprocess.Popen(command.split())
        return p.pid

examples = """
EXAMPLES:

./memleak -p $(pidof allocs)
        Trace allocations and display a summary of "leaked" (outstanding)
        allocations every 5 seconds
./memleak -p $(pidof allocs) -t
        Trace allocations and display each individual allocator function call
./memleak -ap $(pidof allocs) 10
        Trace allocations and display allocated addresses, sizes, and stacks
        every 10 seconds for outstanding allocations
./memleak -c "./allocs"
        Run the specified command and trace its allocations
./memleak
        Trace allocations in kernel mode and display a summary of outstanding
        allocations every 5 seconds
./memleak -o 60000
        Trace allocations in kernel mode and display a summary of outstanding
        allocations that are at least one minute (60 seconds) old
./memleak -s 5
        Trace roughly every 5th allocation, to reduce overhead
"""

description = """
Trace outstanding memory allocations that weren't freed.
Supports both user-mode allocations made with libc functions and kernel-mode
allocations made with kmalloc/kmem_cache_alloc/get_free_pages and corresponding
memory release functions.
"""

parser = argparse.ArgumentParser(description=description,
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=examples)
parser.add_argument("-p", "--pid", type=int, default=-1,
        help="the PID to trace; if not specified, trace kernel allocs")
parser.add_argument("-t", "--trace", action="store_true",
        help="print trace messages for each alloc/free call")
parser.add_argument("interval", nargs="?", default=5, type=int,
        help="interval in seconds to print outstanding allocations")
parser.add_argument("count", nargs="?", type=int,
        help="number of times to print the report before exiting")
parser.add_argument("-a", "--show-allocs", default=False, action="store_true",
        help="show allocation addresses and sizes as well as call stacks")
parser.add_argument("-o", "--older", default=500, type=int,
        help="prune allocations younger than this age in milliseconds")
parser.add_argument("-c", "--command",
        help="execute and trace the specified command")
parser.add_argument("--combined-only", default=False, action="store_true",
        help="show combined allocation statistics only")
parser.add_argument("--wa-missing-free", default=False, action="store_true",
        help="Workaround to alleviate misjudgments when free is missing")
parser.add_argument("-s", "--sample-rate", default=1, type=int,
        help="sample every N-th allocation to decrease the overhead")
parser.add_argument("-T", "--top", type=int, default=10,
        help="display only this many top allocating stacks (by size)")
parser.add_argument("-z", "--min-size", type=int,
        help="capture only allocations larger than this size")
parser.add_argument("-Z", "--max-size", type=int,
        help="capture only allocations smaller than this size")
parser.add_argument("-O", "--obj", type=str, default="c",
        help="attach to allocator functions in the specified object")
parser.add_argument("--ebpf", action="store_true",
        help=argparse.SUPPRESS)
parser.add_argument("--percpu", default=False, action="store_true",
        help="trace percpu allocations")
parser.add_argument("--lib", default=False, action="store_true",
        help="collect by lib")
parser.add_argument("-C", "--comm",
        help="trace the specified command")

args = parser.parse_args()

pid = args.pid
command = args.command or None
kernel_trace = True
trace_all = args.trace
interval = args.interval
min_age_ns = 1e6 * args.older
sample_every_n = args.sample_rate
num_prints = args.count
top_stacks = args.top
min_size = args.min_size
max_size = args.max_size
obj = args.obj
lib = args.lib
comm = args.comm

if min_size is not None and max_size is not None and min_size > max_size:
        print("min_size (-z) can't be greater than max_size (-Z)")
        exit(1)

if command is not None:
        print("Executing '%s' and tracing the resulting process." % command)
        pid = run_command_get_pid(command)


bpf_source = ''

rootPath = os.path.dirname(os.path.abspath(__file__))
with open(os.path.join(rootPath, "memleak.user.c"), 'r') as f:
        bpf_source = f.read()
        bpf_source += '\n'

if kernel_trace:
        if args.percpu:
                with open(os.path.join(rootPath, "memleak.kernel.percpu.c"), 'r') as f:
                        bpf_source += f.read()
        else:
                with open(os.path.join(rootPath, "memleak.kernel.c"), 'r') as f:
                        bpf_source += f.read()

if kernel_trace:
    bpf_source = bpf_source.replace("WORKAROUND_MISSING_FREE", "1"
                                    if args.wa_missing_free else "0")

bpf_source = bpf_source.replace("SHOULD_PRINT", "1" if trace_all else "0")
bpf_source = bpf_source.replace("SAMPLE_EVERY_N", str(sample_every_n))
bpf_source = bpf_source.replace("PAGE_SIZE", str(resource.getpagesize()))

size_filter = ""
if min_size is not None and max_size is not None:
        size_filter = "if (size < %d || size > %d) return 0;" % \
                      (min_size, max_size)
elif min_size is not None:
        size_filter = "if (size < %d) return 0;" % min_size
elif max_size is not None:
        size_filter = "if (size > %d) return 0;" % max_size
bpf_source = bpf_source.replace("SIZE_FILTER", size_filter)

stack_flags = "0"
#if not kernel_trace:
stack_flags += "|BPF_F_USER_STACK"
bpf_source = bpf_source.replace("STACK_FLAGS", stack_flags)
if comm:
        bpf_source = bpf_source.replace("COMM_CARED", comm)
        bpf_source = bpf_source.replace("CARE_ALL", '0')
else:
        bpf_source = bpf_source.replace("CARE_ALL", '1')

if args.ebpf:
    with open(os.path.join(rootPath, "memleak.out.c"), 'w') as f:
           f.write(bpf_source)
    #print(bpf_source)

bpf = BPF(text=bpf_source)

if not kernel_trace:
        print("Attaching to pid %d, Ctrl+C to quit." % -1)

        def attach_probes(sym, fn_prefix=None, can_fail=False):
                if fn_prefix is None:
                        fn_prefix = sym

                try:
                        bpf.attach_uprobe(name=obj, sym=sym,
                                          fn_name=fn_prefix + "_enter",
                                          pid=-1)
                        bpf.attach_uretprobe(name=obj, sym=sym,
                                             fn_name=fn_prefix + "_exit",
                                             pid=-1)
                except Exception:
                        if can_fail:
                                return
                        else:
                                raise

        attach_probes("malloc")
        attach_probes("calloc")
        attach_probes("realloc")
        attach_probes("posix_memalign")
        attach_probes("valloc", can_fail=True) # failed on Android, is deprecated in libc.so from bionic directory
        attach_probes("memalign")
        attach_probes("pvalloc", can_fail=True) # failed on Android, is deprecated in libc.so from bionic directory
        attach_probes("aligned_alloc", can_fail=True)  # added in C11
        bpf.attach_uprobe(name=obj, sym="free", fn_name="free_enter",
                                  pid=-1)

else:
        print("Attaching to kernel allocators, Ctrl+C to quit.")

        # No probe attaching here. Allocations are counted by attaching to
        # tracepoints.
        #
        # Memory allocations in Linux kernel are not limited to malloc/free
        # equivalents. It's also common to allocate a memory page or multiple
        # pages. Page allocator have two interfaces, one working with page
        # frame numbers (PFN), while other working with page addresses. It's
        # possible to allocate pages with one kind of functions, and free them
        # with another. Code in kernel can easy convert PFNs to addresses and
        # back, but it's hard to do the same in eBPF kprobe without fragile
        # hacks.
        #
        # Fortunately, Linux exposes tracepoints for memory allocations, which
        # can be instrumented by eBPF programs. Tracepoint for page allocations
        # gives access to PFNs for both allocator interfaces. So there is no
        # need to guess which allocation corresponds to which free.

def print_outstanding():
        print("[%s] Top %d stacks with outstanding allocations:" %
              (datetime.now().strftime("%H:%M:%S"), top_stacks))
        alloc_info = {}
        allocs = bpf["allocs"]
        stack_traces = bpf["stack_traces"]
        for address, info in sorted(allocs.items(), key=lambda a: a[1].size):
                # if BPF.monotonic_time() - min_age_ns < info.timestamp_ns:
                #         continue
                if info.stack_id < 0:
                        continue
                if info.stack_id in alloc_info:
                        alloc_info[info.stack_id].update(info.size)
                else:
                        stack = list(stack_traces.walk(info.stack_id))
                        combined = []
                        for addr in stack:
                                combined.append(bpf.sym(addr, pid,
                                        show_module=True, show_offset=True))
                        alloc_info[info.stack_id] = Allocation(combined,
                                                               info.size)
                if args.show_allocs:
                        print("\taddr = %x size = %s" %
                              (address.value, info.size))
        to_show = sorted(alloc_info.values(),
                         key=lambda a: a.size)[-top_stacks:]
        for alloc in to_show:
                print("\t%d bytes in %d allocations from stack\n\t\t%s" %
                      (alloc.size, alloc.count,
                       b"\n\t\t".join(alloc.stack).decode("ascii")))

def print_outstanding_combined():
        stack_traces = bpf["stack_traces"]
        stacks = sorted(bpf["combined_allocs"].items(),
                        key=lambda a: -a[1].total_size)
        cnt = 1
        entries = []
        for stack_id, info in stacks:
                try:
                        trace = []
                        for addr in stack_traces.walk(stack_id.value):
                                sym = bpf.sym(addr, pid,
                                                      show_module=True,
                                                      show_offset=True)
                                trace.append(sym)
                        trace = b"\n\t\t".join(trace).decode("ascii")
                except KeyError:
                        trace = "stack information lost"

                entry = ("\t%d bytes in %d allocations from stack\n\t\t%s" %
                         (info.total_size, info.number_of_allocs, trace))
                entries.append(entry)

                cnt += 1
                if cnt > top_stacks:
                        break

        print("[%s] Top %d stacks with outstanding allocations:" %
              (datetime.now().strftime("%H:%M:%S"), top_stacks))

        print('\n'.join(reversed(entries)))


d3fileIndex = 0
def print_d4_combined():
        global d3fileIndex
        d3fileIndex +=1

        allBytes = 0
        if c_uint32(pid) in bpf["stat"]:
                allBytes = bpf["stat"][c_uint32(pid)]
        print(f'{d3fileIndex} alloc size in bytes:{allBytes}')

        gf = open(os.path.join(rootPath, f"memleak.{d3fileIndex}gf.txt"), 'w')
      
        stack_traces = bpf["stack_traces"]
        stacks = sorted(bpf["combined_allocs"].items(),
                        key=lambda a: -a[1].total_size)

        for stack_id, info in stacks:
                taskcomm = info.comm.decode('ascii')
                if comm and taskcomm != comm:
                        print('Error: task command cared fault: ', taskcomm)
                        continue
                try:
                        trace = []

                        # kernel stack
                        for addr in stack_traces.walk(stack_id.value):
                                sym = bpf.sym(addr, pid,
                                        show_module=True,
                                        show_offset=True)
                                sym = str(sym)
                                trace.append(sym)
                        
                        for addr in stack_traces.walk(stack_id.value):
                                sym = bpf.sym(addr, pid,
                                        show_module=True,
                                        show_offset=True)
                                sym = str(sym)
                                trace.append(sym)
                        trace.append(taskcomm)
                        trace.reverse()
                        traces = ';'.join(trace)
                        gf.write(traces)
                        
                except KeyError:
                        print(f"stack information lost: {stack_id.value}")
                        traceback.print_exc()
                        gf.write(str(stack_id.value))
                gf.write(' ')
                gf.write(str(info.total_size))
                gf.write('\n')
        gf.close()
        print("end")

def print_lib_combined():
        global d3fileIndex
        d3fileIndex +=1

        allBytes = 0
        if c_uint32(pid) in bpf["stat"]:
                allBytes = bpf["stat"][c_uint32(pid)]
        print(f'{d3fileIndex} alloc size in bytes:{allBytes}')

        gf = open(os.path.join(rootPath, f"memleak.{d3fileIndex}lib.txt"), 'w')
      
        stack_traces = bpf["stack_traces"]
        stacks = sorted(bpf["combined_allocs"].items(),
                        key=lambda a: -a[1].total_size)

        for stack_id, info in stacks:
                taskcomm = info.comm.decode('ascii')
                if comm and taskcomm != comm:
                        continue
                try:
                        trace = []
                        for addr in stack_traces.walk(stack_id.value):
                                sym = bpf.sym(addr, pid,
                                        show_module=True,
                                        show_offset=True)
                                #print(f'sym {sym}')
                                s = sym.rfind(b'[')
                                e = sym.rfind(b']')
                                lib = sym[s+1:e].decode('ascii')
                                if lib == '':
                                        print(f'wann lib is emplty: {sym}')
                                if lib not in trace:
                                        trace.append(lib)
                        trace.reverse()
                        traces = ';'.join(trace)
                        gf.write(traces)
                        
                except KeyError:
                        print(f"stack information lost: {stack_id.value}")
                        traceback.print_exc()
                        gf.write(str(stack_id.value))
                gf.write(' ')
                gf.write(str(info.total_size))
                gf.write('\n')
        gf.close()
        print("end")

count_so_far = 0
while True:
        if trace_all:
                print(bpf.trace_fields())
        else:
                # print(bpf.trace_fields())
                # continue;
                try:
                        sleep(interval)
                except KeyboardInterrupt:
                        exit()
                p = f'pidof hi' |p
                stdout,_ = p.communicate()
                if stdout != '':
                        pid = int(stdout)
                        print('wann', 'get pid', pid)
                else:
                        print('wann', 'not catch ', comm, ' continue wait')
                        continue
                if args.combined_only:
                        if lib:
                                print_lib_combined()
                        else:
                                print_d4_combined()
                        #print_outstanding_combined()
                else:
                        print_outstanding()
                sys.stdout.flush()
                count_so_far += 1
                if num_prints is not None and count_so_far >= num_prints:
                        exit()
